import cv2
import threading, time
import numpy as np
from xgolib import XGO
from PIL import Image,ImageDraw,ImageFont
import xgoscreen.LCD_2inch as LCD_2inch
import os  # 新增导入os模块用于路径操作
from ultralytics import YOLO
class DogController():
    def __init__(self):
        self.dog = XGO("xgolite")
        self.display = LCD_2inch.LCD_2inch()
        self.display.Init() 
        self.display.clear()
        self.splash = Image.new("RGB",(320,240),"black")
        self.display.ShowImage(self.splash)
        self.cap = None
        self.camera_still = False
        self.save_path = "./saved_images"  # 图片保存路径
        # 确保保存路径存在
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
        self.model = YOLO("./best.pt")
    def open_camera(self):
        print("📷 正在打开摄像头...")
        if self.cap is None:
            self.cap = cv2.VideoCapture(0)
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
            # self.cap.set(cv2.CAP_PROP_FPS, 30)
            # print("✅ 摄像头初始化完成 (320x240@30fps)")
    
    def save_current_image(self):
        """保存当前帧图像"""
        if self.cap is not None and self.cap.isOpened():
            success, image = self.cap.read()
            if success:
                # 生成带时间戳的文件名
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                filename = f"dog_capture_{timestamp}.jpg"
                filepath = os.path.join(self.save_path, filename)
                
                # 转换颜色空间并保存
                # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                cv2.imwrite(filepath, image)
                print(f"📸 图像已保存到: {filepath}")
            else:
                print("⚠️ 无法获取当前帧")
        else:
            print("⚠️ 摄像头未开启")
    
    # def camera_mode(self):
    #     frame_count = 0
    #     while self.camera_still:
    #         success, image = self.cap.read()
    #         if not success:
    #             continue
    #         frame_count += 1
    #         # 色彩通道转换 BGR -> RGB
    #         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    #         image = cv2.flip(image, 1)  # 水平翻转
    #         # 显示处理后的图像（可选，如果需要看到检测效果）
    #         pil_image = Image.fromarray(image)
    #         self.display.ShowImage(pil_image)
    #         time.sleep(0.033)  # 约30fps
    
    def camera_mode(self):
        if not self.cap.isOpened():
            print("无法打开摄像头")
            return

        while self.cap is not None and self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                print("读取摄像头失败")
                break

            results = self.model.predict(frame, conf=0.25, verbose=False)

            for r in results:
                if r.boxes is not None:
                    for box in r.boxes:
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        conf = float(box.conf[0])
                        cls_id = int(box.cls[0])
                        label = self.model.names[cls_id]
                        text = f"{label} {conf:.2f}"
                        # 绘制矩形框和标签
                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                        cv2.putText(frame, text, (x1, y1 - 10), 
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                        # 色彩通道转换 BGR -> RGB
                        image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                        image_rgb = cv2.flip(image_rgb, 1)  # 水平翻转
                        pil_image = Image.fromarray(image_rgb)
                        # 显示处理后的图像
                        self.display.ShowImage(pil_image)
                        time.sleep(0.033)  # 控制帧率（大约30fps）

    def xgoCamera(self, switch):
        print(f"🎮 摄像头开关: {'开启' if switch else '关闭'}")
        if switch:
            self.open_camera()
            self.camera_still = True
            t = threading.Thread(target=self.camera_mode)  
            t.start() 
            # 调整摄像头角度
            print("🔧 调整机械狗摄像头角度...")
            self.dog.motor(11, -30)
            self.dog.motor(21, -30)
            print("✅ 摄像头角度调整完成")
        else:
            print("🛑 正在关闭摄像头...")
            self.camera_still = False
            time.sleep(0.5)
            if self.cap:
                self.cap.release()
                self.cap = None
            splash = Image.new("RGB", (320, 240), "black")
            self.display.ShowImage(splash)
            print("✅ 摄像头已关闭")

    def close(self):
        """清理资源"""
        print("🧹 正在清理系统资源...")
        self.camera_still = False
        if self.cap:
            self.cap.release()
            print("✅ 摄像头资源已释放")
        self.display.clear()
        print("✅ 显示屏已清理")
        print("💤 系统关闭完成")

# 使用示例
if __name__ == "__main__":
    print("🤖 绝影机械狗圆柱台和长条检测系统")
    print("=" * 60)
    print("📌 按下Enter键保存当前画面")
    print("=" * 60)
    x = DogController()
    try:
        x.xgoCamera(True)
        # 运行检测 (无限循环，直到用户中断)
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断程序")
        print("🛑 正在安全关闭系统...")
        x.xgoCamera(False)
    finally:
        x.close()
        print("\n🎯 程序运行结束")
        print("=" * 60)